# 绘制分类指标的直方图
from matplotlib import pyplot as plt
from sklearn.metrics import roc_curve, auc
import seaborn as sns

def plot_metrics(metrics, title, filename,color):
    color = 'skyblue' if '0' in title else 'lightgreen'
    plt.figure(figsize=(6, 6))
    plt.bar(metrics.keys(), metrics.values(), color=color)
    plt.title(title)
    plt.ylabel('Score')
    plt.tight_layout()
    plt.savefig(filename)
    plt.show()

# 绘制 ROC 曲线和计算 AUC 值
def plot_roc_curve(y_test, y_pred_prob, filename):
    fpr, tpr, _ = roc_curve(y_test, y_pred_prob)
    roc_auc = auc(fpr, tpr)
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC)')
    plt.legend(loc='lower right')
    plt.tight_layout()
    plt.savefig(filename)
    plt.show()

# 绘制混淆矩阵的热图
def plot_confusion_matrix(conf_matrix, filename):
    plt.figure(figsize=(8, 6))
    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues',
                xticklabels=['No Defect', 'Defect'],
                yticklabels=['No Defect', 'Defect'])
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.savefig(filename)
    plt.show()
